package com.lengxf.kafka.admin;

import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.DescribeTopicsResult;
import org.apache.kafka.clients.admin.NewPartitions;
import org.apache.kafka.clients.admin.TopicDescription;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicCollection;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
@AllArgsConstructor
public class AdminService {

    private final AdminClient adminClient;

    /**
     * 查询当前topic的分区数量
     *
     * @Author Lengxf
     */
    @SneakyThrows
    public Integer topicPartitionCount(String topic) {
        DescribeTopicsResult describeTopicsResult = adminClient.describeTopics(TopicCollection.TopicNameCollection.ofTopicNames(List.of(topic)));
        KafkaFuture<TopicDescription> topicDescriptionKafkaFuture = describeTopicsResult.topicNameValues().get(topic);
        if (topicDescriptionKafkaFuture == null) {
            return 0;
        }
        TopicDescription topicDescription = topicDescriptionKafkaFuture.get();
        return topicDescription.partitions().size();
    }

    /**
     * topic分区扩容
     *
     * @Author Lengxf
     */
    public void partitionIncreaseTo(String topic, int totalCount) {
        if (topicPartitionCount(topic) >= totalCount) {
            throw new RuntimeException("final partitions must be over current counts ...");
        }
        Map<String, NewPartitions> newPartitions = new HashMap<>();
        newPartitions.put(topic, NewPartitions.increaseTo(totalCount));
        adminClient.createPartitions(newPartitions);
    }
}
